// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <limits>

#include <ktest/ktest.h>

#include <nebula/core/memory_pool.h>
#include <turbo/utility/status.h>
#include <nebula/testing/ktest_util.h>

namespace nebula {

class TestMemoryPoolBase : public ::testing::Test {
 public:
  virtual ::nebula::MemoryPool* memory_pool() = 0;

  void TestMemoryTracking() {
    auto pool = memory_pool();

    uint8_t* data;
    const auto old_bytes_allocated = pool->bytes_allocated();
    ASSERT_OK(pool->allocate(100, &data));
    EXPECT_EQ(static_cast<uint64_t>(0), reinterpret_cast<uint64_t>(data) % 64);
    ASSERT_EQ(old_bytes_allocated + 100, pool->bytes_allocated());

    uint8_t* data2;
    ASSERT_OK(pool->allocate(27, &data2));
    EXPECT_EQ(static_cast<uint64_t>(0), reinterpret_cast<uint64_t>(data2) % 64);
    ASSERT_EQ(old_bytes_allocated + 127, pool->bytes_allocated());

    pool->free(data, 100);
    ASSERT_EQ(old_bytes_allocated + 27, pool->bytes_allocated());
    pool->free(data2, 27);
    ASSERT_EQ(old_bytes_allocated, pool->bytes_allocated());
  }

  void TestOOM() {
    auto pool = memory_pool();

    uint8_t* data;
    int64_t max_alloc = std::min<uint64_t>(std::numeric_limits<int64_t>::max(),
                                           std::numeric_limits<size_t>::max());
    // subtract 63 to prevent overflow after the size is aligned
    for (int64_t to_alloc : {max_alloc, max_alloc - 63, max_alloc - 127}) {
      ASSERT_RAISES(turbo::StatusCode::kResourceExhausted, pool->allocate(to_alloc, &data));
    }
  }

  void TestReallocate() {
    auto pool = memory_pool();

    uint8_t* data;
    ASSERT_OK(pool->allocate(10, &data));
    ASSERT_EQ(10, pool->bytes_allocated());
    data[0] = 35;
    data[9] = 12;

    // Expand
    ASSERT_OK(pool->reallocate(10, 20, &data));
    ASSERT_EQ(data[9], 12);
    ASSERT_EQ(20, pool->bytes_allocated());

    // Shrink
    ASSERT_OK(pool->reallocate(20, 5, &data));
    ASSERT_EQ(data[0], 35);
    ASSERT_EQ(5, pool->bytes_allocated());

    // Free
    pool->free(data, 5);
    ASSERT_EQ(0, pool->bytes_allocated());
  }

  void TestAlignment() {
    auto pool = memory_pool();
    {
      uint8_t* data64;
      ASSERT_OK(pool->allocate(10, &data64));
      ASSERT_EQ(reinterpret_cast<uintptr_t>(data64) % kDefaultBufferAlignment, 0);
      pool->free(data64, 10);
    }

    {
      uint8_t* data512;
      ASSERT_OK(pool->allocate(10, 512, &data512));
      ASSERT_EQ(reinterpret_cast<uintptr_t>(data512) % 512, 0);
      pool->free(data512, 10, 512);
    }
  }
};

}  // namespace nebula
